import numpy as np
import h5py as h5
import glob
import os
import json
from tri_mesh2d import *

def collect(groupName, dataSetName = None, itime=None, path="data",prefix="data", edge_name = None):
    file_list = glob.glob(os.path.join(path, prefix+"*.h5"))
    #print("file_list: ", file_list)
    file_list.sort()
    #print("file_list: ", file_list)

    #read easypic2d.config
    config_file = open("data/easypic2d.config")
    json_dic = json.load(config_file)
    pmsh_mesh_name = json_dic["pmsh_mesh_name"]

    #read mesh
    mesh = TriMesh2D()
    mesh.read_pmsh_mesh(pmsh_mesh_name)

    #print("[collect.py::collect] ", groupName, itime)


    if groupName == "fields" or groupName == "vector_fields":
        if itime == None:
            i = 0
            for fileName in file_list:
                f = h5.File(fileName, 'r')
                if dataSetName == None:
                    val = f[groupName]
                else:
                    val = f[groupName][dataSetName][:]
                val = val[np.newaxis, :]
                if i == 0:
                    data = val
                else:
                    data = np.append(data, val, axis = 0)
                i = i + 1
                f.close()
            return data
        else:
                f = h5.File(file_list[itime])
                if dataSetName == None:
                    val = f[groupName]
                else:
                    val = f[groupName][dataSetName][:]
                f.close()
                return val

    elif groupName == "edge_fields":
        if itime == None:
            i = 0
            for fileName in file_list:
                f = h5.File(fileName, 'r')
                if dataSetName == None:
                    val = f[groupName]
                else:
                    val = f[groupName][dataSetName][:]
                val = val[:]
                val = val[np.newaxis, :]
                if i == 0:
                    data = val
                else:
                    data = np.append(data, val, axis = 0)
                i = i + 1
                f.close()
            return data
        else:
            f = h5.File(file_list[itime], 'r')
            if dataSetName == None:
                val = f[groupName][:]
            else:
                val = f[groupName][dataSetName][:]
            edge_field_all = np.zeros((2, val.shape[0]))
            edge_field_all[1,:] = val
            edge_field_all[0,:] = mesh.bar_x
            #edge_field_all[0,:] = mesh.bar_y
            #edge_field_all[0,:] = mesh.bar_distance
            #print("[collect.py::collect] ", val)

            edge = mesh.edge_list[edge_name]

            edge_field = np.zeros((2, edge.shape[0]))
            for i in range(0, edge.shape[0]):
                edge_field[0, i] = edge_field_all[0, edge[i]]
                edge_field[1, i] = edge_field_all[1, edge[i]]

            f.close()
            return edge_field



#pn = collect("/Diagnostic", "particleNumber")     
